package cn.monkey.transport.mq.rabbitmq;

import com.rabbitmq.client.*;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeoutException;

public class DefaultReactiveRabbitMqChannelPool implements ReactiveRabbitMqChannelPool {

    private final ConnectionFactory connectionFactory;

    private final Map<Address, Channel> channelMap;

    public DefaultReactiveRabbitMqChannelPool(ConnectionFactory connectionFactory) {
        connectionFactory.useNio();
        this.connectionFactory = connectionFactory;
        this.channelMap = new ConcurrentHashMap<>();
    }

    public DefaultReactiveRabbitMqChannelPool() {
        this(new ConnectionFactory());
    }


    @Override
    public Mono<Channel> findOrCreate(Address address, Map<String, Object> properties) {
        return Mono.justOrEmpty(this.channelMap.get(address))
                .switchIfEmpty(Mono.defer(() -> Mono.create(sink -> CompletableFuture.runAsync(() -> {
                    try {
                        Connection connection = this.connectionFactory.newConnection(List.of(address));
                        Channel channel = this.channelMap.computeIfAbsent(address, a -> {
                            try {
                                return connection.createChannel();
                            } catch (IOException e) {
                                throw new RuntimeException(e);
                            }
                        });
                        sink.success(channel);
                    } catch (IOException | TimeoutException | RuntimeException e) {
                        if (e instanceof RuntimeException re) {
                            sink.error(re.getCause());
                            return;
                        }
                        sink.error(e);
                    }
                }))));
    }

    @Override
    public Mono<Void> destroy(Channel... channels) {
        Iterator<Map.Entry<Address, Channel>> iterator = this.channelMap.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<Address, Channel> next = iterator.next();
            for (Channel channel : channels) {
                if (next.getValue() == channel) {
                    iterator.remove();
                }
            }
        }
        CompletableFuture<?>[] futures = Arrays.stream(channels)
                .map(c -> CompletableFuture.runAsync(() -> {
                    try {
                        c.close();
                    } catch (IOException | TimeoutException e) {
                        throw new RuntimeException(e);
                    }
                }))
                .toArray(CompletableFuture[]::new);
        return Mono.fromCompletionStage(CompletableFuture.allOf(futures));
    }
}
